package com.northpool.operator.ai.deepseek;

import cn.hutool.ai.AIServiceFactory;
import cn.hutool.ai.ModelName;
import cn.hutool.ai.core.AIConfigBuilder;
import cn.hutool.ai.core.Message;
import cn.hutool.ai.model.deepseek.DeepSeekService;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.northpool.operator.annotation.InputParam;
import com.northpool.operator.annotation.InputParams;
import com.northpool.operator.annotation.OperatorInfo;
import com.northpool.operator.annotation.OutputParam;
import com.northpool.operator.annotation.OutputParams;
import com.northpool.operator.annotation.ParamEnum;
import com.northpool.operator.annotation.ParamEnums;
import com.northpool.operator.common.algorithm.OperatorTestBuild;
import com.northpool.operator.common.algorithm.TestParamBuilder;
import com.northpool.operator.param.ParamTool;
import com.northpool.operator.type.TypeEnum;
import com.northpool.operator.utils.StpOperatorHelper;

import java.util.LinkedList;
import java.util.List;

@OperatorInfo(
        ref = "DeepSeekChatOperator",
        name = "AI大模型",
        version = "deepseek",
        description = "AI大模型单次对话",
        enums = {
                @ParamEnums(ref = "model", value = {
                        @ParamEnum(ref = "deepseek-chat", name = "聊天模型", description = "deepseek-chat 模型已经升级为 DeepSeek-V3"),
                        @ParamEnum(ref = "deepseek-reasoner", name = "推理模型", description = "deepseek-reasoner 模型为新模型 DeepSeek-R1"),
                })
        }
)
@OutputParams({
        @OutputParam(ref = "content", name = "AI的回答", desc = "AI的回答")
})
public class DeepSeekChatOperator extends StpOperatorHelper {

    @InputParams({
            @InputParam(ref = "question", name = "问题", desc = "问题描述", required = true),
            @InputParam(ref = "apiKey", name = "API Key", desc = "DeepSeek大模型服务的API Key", required = true),
            @InputParam(ref = "model", name = "模型版本", desc = "模型版本，分为聊天模型和推理模型", required = true, type = TypeEnum.ENUM, enumRef = "model", value = "deepseek-chat"),
            @InputParam(ref = "roleDesc", name = "角色描述", desc = "大模型的角色描述, 如: 你是一个中英文翻译专家，将用户输入的中文翻译成英文，或将用户输入的英文翻译成中文。"),
    })
    public DeepSeekChatOperator() {
    }

    public static void main(String[] args) {
        OperatorTestBuild.execute(DeepSeekChatOperator.class, TestParamBuilder.aBuilder()
                .param("question", "写一首七言绝句")
                .param("apiKey", "sk-4da549f643924d9eb237587e9dd0f739")
                .param("roleDesc", "")
                .build()
        );
    }

    @Override
    public void execute(ParamTool paramTool) {
        String key = paramTool.getStr("apiKey");
        String model = paramTool.getStr("model");
        String roleDesc = paramTool.getStr("roleDesc");

        DeepSeekService deepSeekService = AIServiceFactory.getAIService(
                new AIConfigBuilder(ModelName.DEEPSEEK
                        .getValue())
                        .setModel(model)
                        .setApiKey(key)
                        .build(),
                DeepSeekService.class);

        String question = paramTool.getStr("question");
        List<Message> messageList = new LinkedList<>();
        if (StrUtil.isNotBlank(roleDesc)) {
            messageList.add(new Message("system", roleDesc));
        }
        messageList.add(new Message("user", question));
        String response = deepSeekService.chat(messageList);

        JSONObject entries = JSONUtil.parseObj(response);
        JSONArray choices = entries.getJSONArray("choices");
        JSONObject choice = (JSONObject) choices.stream().findFirst().orElseThrow(() -> new RuntimeException("回答为空"));

        JSONObject message = choice.getJSONObject("message");
        String content = message.getStr("content");

        this.result("content", content);
    }
}
